import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import time

transform = transforms.Compose([transforms.ToTensor()])
train_datasets = datasets.MNIST(root="../data", train=True, download=True, transform=transform)
test_datasets = datasets.MNIST(root="../data", train=False, download=True, transform=transform)
train_loader = DataLoader(dataset=train_datasets, batch_size=64, shuffle=True, num_workers=2)
test_loader = DataLoader(dataset=test_datasets, batch_size=64, shuffle=True, num_workers=2)


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 10, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(10, 20, kernel_size=5),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(320, 50),
            torch.nn.Linear(50, 10)
        )

    def forward(self, x):
        batch_size = x.size(0)
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(batch_size, -1)
        x = self.fc(x)
        return x


device = "cuda:0" if torch.cuda.is_available() else "cpu"
model = Net()
model.to(device)

loss = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params=model.parameters(), lr=0.01)


def train(epoch):
    loss_total = 0.0
    total = 0
    correct = 0
    for batch_index, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.to(device)
        outputs = model(inputs)
        optimizer.zero_grad()
        loss_val = loss(outputs, targets)
        loss_val.backward()
        optimizer.step()
        loss_total += loss_val.item()
        total += inputs.size(0)
        _, predicted = torch.max(outputs.data, dim=1)
        correct += (predicted == targets).sum().item()
        if batch_index % 300 == 0:
            print(
                "[%d, %5d], loss: %.2f, acc: %.2f %%" % (epoch + 1, batch_index + 1, loss_total, 100 * correct / total))
            loss_total = 0.0
            total = 0
            correct = 0


def test(epoch):
    total = 0
    correct = 0
    with torch.no_grad():
        for batch_index, (inputs, targets) in enumerate(test_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, dim=1)
            total += inputs.size(0)
            correct += (predicted == targets).sum().item()
        acc = 100 * correct / total
        print("acc: %.2f" % acc)


if __name__ == "__main__":
    start_time = time.time()
    for epoch in range(10):
        train(epoch)
        test(epoch)
    end_time = time.time()
    print("time: %.2f " % ((end_time - start_time) / 60))
